package org.example.rabbitmq.consumer.rate;

import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.data.redis.core.StringRedisTemplate;

import java.time.Duration;
import java.util.concurrent.TimeUnit;

/**
 * 队列消费限流器
 * <p>
 * 可以用于多进程，可以避免单点问题
 * 除了Redis，也可以使用Zookeeper等中间件
 */
public class RedisRateLimiter implements RateLimiter {
    private static final String MQ_CONSUMER_RATE_PREFIX = "mq_rate_";
    private int durationSeconds = 60;
    private int tickMilliSeconds = 1000;
    private int retryTimes = 100;

    public RedisRateLimiter(int durationSeconds) {
        this.durationSeconds = durationSeconds;
    }

    @Autowired
    StringRedisTemplate redisTemplate;

    @Override
    public void acquire(String queueName) throws InterruptedException {
        boolean created = setKeyNx(queueName);
        if (created) {
            return;
        }
        // 循环方式获取是示例代码，等待时间和请求次数都不容易确定，最好使用watcher通知机制
        int i = 0;
        for (; i < retryTimes; i++) {
            created = setKeyNx(queueName);
            if (b(created)) {
                return;
            } else {
                TimeUnit.MILLISECONDS.sleep(tickMilliSeconds);
            }
        }
        if (retryTimes == i) {
            // timeout, the message should be requeue
            throw new RateLimiterTimeoutException("acquire lock timeout, can not acquire lock for queue " + queueName + " in " + retryTimes * tickMilliSeconds + "ms");
        }
    }

    @Override
    public void release(String queueName) {
        redisTemplate.delete(getKey(queueName));
    }

    @Override
    public void addQueueRateLimiter(String queueName) {
        // do nothing
    }

    private boolean setKeyNx(String queueName) {
        // 这里value可以设置为 当前时间戳，释放key时可以统计一共锁定了多少时间
        Boolean created = redisTemplate.opsForValue().setIfAbsent(getKey(queueName), queueName, Duration.ofSeconds(durationSeconds));
        return b(created);
    }

    private String getKey(String queueName) {
        // 如果可以允许同时由两个执行，则需要传入一个和Message上下文相关的值组成键值
        return MQ_CONSUMER_RATE_PREFIX + queueName;
    }

    private boolean b(Boolean b) {
        return b != null && b.booleanValue();
    }
}
